//#include "StdAfx.h"
#include <strstream>

#include "Test_BayesianFilter.h"
#include "BayesianFilter.h"
#include "Log.h"

#include <iostream>
#include <algorithm>
#include <math.h>

#define PROBABILITY_THRESHOLD 0.0

__BayesianProbability::__BayesianProbability() : trueCount(0), falseCount(0)
{
}

__BayesianProbability __BayesianProbability::operator = (const __BayesianProbability& rval)
{
	trueCount = rval.trueCount;
	falseCount = rval.falseCount;
	return *this;
}

BayesianFilter::BayesianFilter(void) : ProbabilityThreshold(PROBABILITY_THRESHOLD)
{
}

BayesianFilter::~BayesianFilter(void)
{
	map<string,__BayesianTag*>::iterator iterTag;
	for ( iterTag = tags.begin () ; iterTag != tags.end() ; iterTag++ )
		delete iterTag->second;

	map<string,__BayesianProbability*>::iterator iterWord;
	for ( iterWord = words.begin() ; iterWord != words.end() ; iterWord++ )
		delete iterWord->second;
}

void BayesianFilter::AddTag(string tagName)
{
	LOG1("  BayesianFilter::AddTag");
	lock.Enter ();

	map<string,__BayesianTag*>::iterator iterTag = tags.find(tagName);
	if ( iterTag != tags.end () )
	{
		lock.Release ();
		throw new TagAlreadyExistsException ();
	}

	struct __BayesianTag* newTag = new struct __BayesianTag;
	newTag->pos = (unsigned long) tags.size();
	newTag->wordCount = 0;
	newTag->trueCount = 0;
	newTag->falseCount = 0;

	int oldSize = (int) tags.size ();
	tags.insert ( make_pair ( tagName , newTag ) );
	  
	map<string,__BayesianProbability*>::iterator iterWord;
	map<string,__BayesianProbability*> tmpMap;
	for ( iterWord = words.begin () ; iterWord != words.end() ; iterWord++ )
	{
		__BayesianProbability* newProbabilities = new __BayesianProbability [ oldSize + 1 ];
		newProbabilities[oldSize].trueCount = 0;
		newProbabilities[oldSize].falseCount = 0;
		for ( int i = 0 ; i < oldSize ; i++ )
		{
			newProbabilities[i].trueCount = iterWord->second[i].trueCount;
			newProbabilities[i].falseCount = iterWord->second[i].falseCount;
		}

		delete [] iterWord->second;
		words[iterWord->first] = newProbabilities;
	}

	lock.Release ();
}

void BayesianFilter::RemoveTag(std::string tagName)
{
	LOG1("  BayesianFilter::RemoveTag");
	lock.Enter ();

	map<string,__BayesianTag*>::iterator iterTag = tags.find(tagName);
	if ( iterTag == tags.end () )
	{
		lock.Release ();
		throw new TagNotExistsException ();
	}

	__BayesianTag* toDie = iterTag->second;
	tags.erase(iterTag->first);

	for ( iterTag = tags.begin() ; iterTag != tags.end() ; iterTag++ )
		if ( iterTag->second->pos > toDie->pos )
			iterTag->second->pos--;

	unsigned int newSize = (unsigned int) tags.size ();
	map<string,__BayesianProbability*>::iterator iterWord;
	for ( iterWord = words.begin () ; iterWord != words.end() ; iterWord++ )
	{
		__BayesianProbability* newProbabilities = new __BayesianProbability [ newSize ];

		for ( unsigned int i = 0 ; i < toDie->pos ; i++ )
			newProbabilities[i] = iterWord->second[i];

		for ( unsigned int i = toDie->pos + 1 ; i <= newSize ; i++ )
			newProbabilities[i-1] = iterWord->second[i];		

		delete [] iterWord->second;
		words[iterWord->first] = newProbabilities;
	}

	delete toDie;
	lock.Release ();
}

/*
 * unsigned long* FindTagPos ( set<string>& )
 *
 *   Devuelve un array de unsigned long, del mismo tama~no que el conjunto
 * pasado, con los nombres de tag traducidos a posiciones. Notese que el 
 * array queda ordenado para acelerar luego la actualizacion de las filas.
 *
 */

unsigned long* BayesianFilter::FindTagPos(set<string> &trainTags)
{
	int size = (int) trainTags.size ();
	unsigned long* tagPos = new unsigned long [ size ];

	set<string>::iterator iterTags;
	int i = 0;
	for ( iterTags = trainTags.begin() ; iterTags != trainTags.end() ; iterTags++, i++ )
	{
		map<string,__BayesianTag*>::iterator tagFound = tags.find(*iterTags);
		if ( tagFound == tags.end () )
		{
			delete [] tagPos;
			lock.Release ();
			throw new TagNotExistsException ();
		}

		tagPos[i] = tagFound->second->pos;
	}

	sort ( tagPos , tagPos + size );
	return tagPos;
}

/*
 * __BayesianProbability* ProbabilityRow ( string )
 *
 *   Retorna el puntero a la fila de probabilidades para todos los tags
 * de la palabra indicada. Siempre retorna una fila valida, puesto que no
 * existe actualmente la palabra en el diccionario, la agrega e inicializa
 * la columna de forma adecuada.
 *
 */

__BayesianProbability* BayesianFilter::ProbabilityRow ( string word )
{
	__BayesianProbability* retval;
	map<string,__BayesianProbability*>::iterator iter;
	iter = words.find(word);

	if ( iter != words.end() )
		retval = iter->second;
	else
	{
		int size = (int) tags.size ();
		retval = new __BayesianProbability [ size ];
		for ( int i = 0 ; i < size ; i++ )
		{
			retval[i].trueCount = 0;
			retval[i].falseCount = 0;
		}

		words[word] = retval;
	}

	return retval;
}

/*
 * UpdateRow ( __BayesianProbability , unsigned long* , int )
 *
 *   Actualiza la fila de probabilidades indicada, incrementando el conteo
 * de palabra vista en cada una de las categorias de array pasado e incrementando
 * el conteo de palabra vista en la negacion de cada una de las categorias no
 * incluidas en dicho array. 
 *
 */

void BayesianFilter::UpdateRow ( __BayesianProbability* p, unsigned long* tagPos, int size , UpdateModes updateMode )
{
	int rowSize = (int) tags.size ();
	int rowTag = 0;
	int updatePos = 0;
	int dump = 0;

	for ( int i = 0 ; i < rowSize ; i++ )
	{
		if ( dump || i != tagPos [ updatePos ] )
		{
			if ( updateMode == TrainMode )
				p[i].falseCount++;
		}
		else
		{
			if ( updateMode == TrainMode ) 
				p[i].trueCount++;
			else if ( updateMode == AddMode )
			{
				p[i].trueCount++;
				if ( p[i].falseCount != 0 )
					p[i].falseCount--;
			}
			else if ( updateMode == RemoveMode )
			{
				p[i].falseCount++;
				if ( p[i].trueCount != 0 )
					p[i].trueCount--;
			}

			if ( updatePos++ == size ) dump = 1;
		}
	}
}

void BayesianFilter::Train ( set<string>& trainWords, set<string>& trainTags , UpdateModes updateMode )
{
	if ( trainTags.empty () )
		throw new EmptyTagSet ();

#ifdef DEBUG_BAYES
	LOG1("BayesianFilter::Train");
	if ( updateMode == AddMode )
		LOG1("  Mode: Add");
	else if ( updateMode == RemoveMode )
		LOG1("  Mode: Remove");
	
	string __tags;
	set<string>::iterator __iterStr;
	for ( __iterStr = trainTags.begin() ; __iterStr != trainTags.end() ; __iterStr++ )
		__tags += *__iterStr + "; ";
	
	LOG2("  Tags: ",__tags);
	
	string __words;
	for ( __iterStr = trainWords.begin() ; __iterStr != trainWords.end() ; __iterStr++ )
		__words += *__iterStr + "; ";

	LOG2("  Words: ",__words);
#endif

	lock.Enter ();
	int size = (int) tags.size ();
	unsigned long* tagPos = FindTagPos ( trainTags );

#ifdef DEBUG_BAYES
	string __tagPos;
	strstream __tmps;
	string __tmp;
	for ( int __i = 0 ; __i < size ; __i++ )
	{	__tmps << tagPos[__i];
		__tmps >> __tmp;
		__tagPos += __tmp + ";";
	}

	LOG2("  Tag Pos: ",__tagPos);
	LOG1("  Training...");
#endif

	set<string>::iterator iter;
	for ( iter = trainWords.begin() ; iter != trainWords.end() ; iter++ )
	{
		__BayesianProbability* p = ProbabilityRow ( *iter );
		UpdateRow ( p , tagPos , size , updateMode );
	}

	int wc = (int) trainWords.size ();
	map<string,__BayesianTag*>::iterator iterTags;
	for ( iterTags = tags.begin() ; iterTags != tags.end() ; iterTags++ )
	{
		iterTags->second->wordCount += wc;
		if ( trainTags.find(iterTags->first) != trainTags.end() )
		{
			if 	( updateMode == TrainMode )
				iterTags->second->trueCount++;
			else if ( updateMode == AddMode )
			{
				iterTags->second->trueCount++;
				if ( iterTags->second->falseCount != 0 )
					iterTags->second->falseCount--;
			}
			else if ( updateMode == RemoveMode )
			{
				iterTags->second->falseCount++;
				if ( iterTags->second->trueCount != 0 )
					iterTags->second->trueCount--;
			}
		}
		else
		{
			if ( updateMode == TrainMode )
				iterTags->second->falseCount++;
		}
	}

	delete [] tagPos;
	lock.Release();
}

void BayesianFilter::Train (set<string>& words, set<string>& tags)
{
	Train ( words , tags , TrainMode );
}

void BayesianFilter::Retrain ( set<string>& words , set<string>& removeTags , set<string>& addTags )
{
	try
	{
		Train ( words , removeTags , RemoveMode );
	}
	catch ( EmptyTagSet* ex )
	{
		// No es un error no querer quitar ninguna categoria
		delete ex;
	}

	try
	{
		Train ( words , addTags , AddMode );
	}
	catch ( EmptyTagSet* ex )
	{
		// No es un error no querer agregar ninguna categoria
		delete ex;
	}
}

set<BayesianFilterResult*> BayesianFilter::Classify(std::set<string> &words)
{
#ifdef DEBUG_BAYES
	LOG1("BayesianFilter::Classify");
	
	string __words;
	set<string>::iterator __iterStr;
	for ( __iterStr = words.begin() ; __iterStr != words.end() ; __iterStr++ )
		__words += *__iterStr + "; ";

	LOG2("  Words: ",__words);
#endif
	
	lock.Enter ();
	int tagPos;
	int wordCount = (int) words.size ();
	int tagSize = (int) tags.size ();
	unsigned long* hr = new unsigned long [ tagSize ];
	double* p = new double [ tagSize ];

#ifdef DEBUG_BAYES
	LOG2("  Word Count: ",wordCount);
	LOG2("  Tag Size:   ",tagSize);
	LOG1("  Adding a priori probability...");
#endif

	map<string,__BayesianTag*>::iterator tagIter;
	for ( tagIter = tags.begin() ; tagIter != tags.end() ; tagIter++, tagPos++ ) 
	{
		tagPos = tagIter->second->pos;
		hr[tagPos] = 0;

		// Sumarle uno al conteo de TAG y NO TAG es la manera mas elegante
		// de resolver simultaneamente los problemas de division por cero
		// y probabilidades nulas.
		p[tagPos] = log ( (double) (tagIter->second->trueCount + 1) / 
						  (double) (tagIter->second->falseCount + 1) );

#ifdef DEBUG_BAYES
		LOG4("    Pos: ",tagPos,"  p: ",p[tagPos]);
#endif
	}

#ifdef DEBUG_BAYES
	LOG1("  Visiting words...");
#endif
	
	set<string>::iterator iter;
	for ( iter = words.begin() ; iter != words.end() ; iter++ )
	{
#ifdef DEBUG_BAYES		
		LOG2("    Word: ",*iter);
#endif
		__BayesianProbability* r = ProbabilityRow ( *iter );
		for ( int j = 0 ; j < tagSize ; j++ )
		{
			// Ver comentario anterior
			p[j] += log ( (double) (r[j].trueCount + 1) / 
				          (double) (r[j].falseCount + 1) );

			if ( r[j].trueCount != 0 || r[j].falseCount != 0 )
				hr[j]++;

#ifdef DEBUG_BAYES
			LOG10("      (",j,") TrueCount: ",r[j].trueCount," FalseCount: ",
				r[j].falseCount," p: ",p[j]," hr: ",hr[j]);
#endif
		}
	}

#ifdef DEBUG_BAYES
	LOG1("  Generating final result...");	
#endif
	
	set<BayesianFilterResult*> retval;
	for ( tagIter = tags.begin() ; tagIter != tags.end() ; tagIter++, tagPos++ )
	{
		tagPos = tagIter->second->pos;

		//if ( p[tagPos] >= ProbabilityThreshold )
		//{
			struct BayesianFilterResult* tmp = new struct BayesianFilterResult ();
			tmp->Tag = tagIter->first;
			tmp->Witness  = p [tagPos];
			tmp->HitRatio = (double) hr[tagPos] / (double) wordCount;

			retval.insert(tmp);
		//}
#ifdef DEBUG_BAYES
		LOG6("  Tag: ",tmp->Tag," Witness: ",tmp->Witness," HitRatio: ",tmp->HitRatio);
#endif
	}

	lock.Release ();
	delete [] hr;
	delete [] p;
	return retval;
}
